package edu.usc.epigenome.dnaase.script;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;

import net.sf.samtools.SAMFileReader;
import net.sf.samtools.SAMRecord;
import net.sf.samtools.util.CloseableIterator;

import org.kohsuke.args4j.Argument;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;
import org.kohsuke.args4j.Option;

import edu.usc.epigenome.genomeLibs.PicardUtils;

public class SamFilter {

	/**
	 * @param args
	 */
		
		//final private static String prefix = "methylCGsRich_ASM_AllSnp_";
		final private static String USAGE = "SamFilter [opts] snpFileName bamFile ...";

		final private static int PURGE_INTERVAL = 20000; // We purge our stored Cpgs once we get this many bases past them.
		//final private static int ALLELE_GA_NUMBER = 1; //for each reads, when there are more than 1 GA position is different from reference sequence, we define it belongs to another allele. 
		
		/**
		 * @param args
		 */
		
		@Option(name="-minMapQ",usage="minimum mapping quality (default 30)")
		protected int minMapQ = 30;
		@Option(name="-minReadCov",usage="minimum read coverage (default 4)")
		protected int minReadCov = 1;
		@Option(name="-minBaseQual",usage="minimum Base quality (default 0)")
		protected int minBaseQual = 0;
		//for cytosine
		@Option(name="-minCTReadCov",usage="minimum CT read coverage (default 3)")
		protected int minCTReadCov = 3;
		
		@Option(name="-debug",usage=" Debugging statements (default false)")
		protected boolean debug = false;
		// receives other command line parameters than options
		@Argument
		private List<String> stringArgs = new ArrayList<String>();

		
		
		/**
		 * @param args
		 */
		public static void main(String[] args)
		throws Exception
		{
			SamFilter test = new SamFilter();
			test.doMain(args);
		}

		public void doMain(String[] args)
		throws Exception {

			CmdLineParser parser = new CmdLineParser(this);
			// if you have a wider console, you could increase the value;
			// here 80 is also the default
			parser.setUsageWidth(80);
			try
			{
				parser.parseArgument(args);
				if (stringArgs.size() < 2) throw new CmdLineException(USAGE);
			
			}
			catch (CmdLineException e)
			{
				System.err.println(e.getMessage());
				// print the list of available options
				parser.printUsage(System.err);
				System.err.println();
				return;
			}

			
			String snpFileName = stringArgs.remove(0);
			String bamFileName = stringArgs.remove(0);
			BufferedReader br = new BufferedReader(new FileReader(snpFileName));
			String line;
			
			int snpHeterLoci = 0;
			int snpHomoLoci = 0;
			int snpHeterLociFilter = 0;
			int snpHomoLociFilter = 0;
			
			int ctSnpNum = 0;
			int nonCtSnpNum = 0;
			int ctSnpNumFilter = 0;
			int nonCtSnpNumFilter = 0;
			
			int snpHeterLociC = 0;
			int snpHomoLociC = 0;
			int snpHeterLociCFilter = 0;
			int snpHomoLociCFilter = 0;
			
			String outputFileName = snpFileName + ".minReadCov_" + minReadCov + ".snpcaller.txt";
			String outputBedFileName = snpFileName + ".minReadCov_" + minReadCov + ".snpcaller.bed";
			String outputCTsnpFileName = snpFileName + ".minReadCov_" + minReadCov + ".snpcaller.ctsnp.txt";
			String outputCTsnpBedFileName = snpFileName + ".minReadCov_" + minReadCov + ".snpcaller.ctsnp.bed";
			String outputNonCTsnpFileName = snpFileName + ".minReadCov_" + minReadCov + ".snpcaller.nonctsnp.txt";
			String outputNonCTsnpBedFileName = snpFileName + ".minReadCov_" + minReadCov + ".snpcaller.nonctsnp.bed";
			
			String cytosineFileName = snpFileName + ".minReadCov_" + minReadCov + ".minCTReadCov_" + minCTReadCov + ".cytosinecaller.txt";
			String cytosineBedFileName = snpFileName + ".minReadCov_" + minReadCov + ".minCTReadCov_" + minCTReadCov + ".cytosinecaller.bed";
			
			
			File inputSamOrBamFile = new File(bamFileName);
			SAMFileReader inputSam = new SAMFileReader(inputSamOrBamFile);
			inputSam.setValidationStringency(SAMFileReader.ValidationStringency.SILENT);
			
			PrintWriter outWriter = new PrintWriter(new File(outputFileName));
			PrintWriter outCytosineWriter = new PrintWriter(new File(cytosineFileName));
			
			PrintWriter outCTsnpWriter = new PrintWriter(new File(outputCTsnpFileName));
			PrintWriter outCTsnpBedWriter = new PrintWriter(new File(outputCTsnpBedFileName));
			
			PrintWriter outNonCTsnpWriter = new PrintWriter(new File(outputNonCTsnpFileName));
			PrintWriter outNonCTsnpBedWriter = new PrintWriter(new File(outputNonCTsnpBedFileName));
			
			PrintWriter outBedWriter = new PrintWriter(new File(outputBedFileName));
			PrintWriter outCytosineBedWriter = new PrintWriter(new File(cytosineBedFileName));
			
			while( (line = br.readLine()) != null){
				String[] tmpArray = line.split("\t");
				String chr = "chr" + tmpArray[1];
				
				
				
				Integer snpPosition =(int) Math.round(Double.parseDouble(tmpArray[2]));
				//push(@out,"$chr\t$start\t$end\t$name\t0\t+\n");
				int start = snpPosition - 1;
				int end = snpPosition;
				String name = tmpArray[0];
				
				boolean snpHomo = tmpArray[3].equalsIgnoreCase(tmpArray[4]);
				
				int totalNum = 0;

				int totalNumC = 0;
				int totalNumT = 0;
				
				if(snpHomo){
					snpHomoLoci++;
					
				}
				else{
					snpHeterLoci++;
					if((tmpArray[5].equalsIgnoreCase("C") && tmpArray[6].equalsIgnoreCase("T")) || (tmpArray[5].equalsIgnoreCase("T") && tmpArray[6].equalsIgnoreCase("C")) || (tmpArray[5].equalsIgnoreCase("A") && tmpArray[6].equalsIgnoreCase("G")) || (tmpArray[5].equalsIgnoreCase("G") && tmpArray[6].equalsIgnoreCase("A"))){
						ctSnpNum++;
					}
					else{
						nonCtSnpNum++;
					}
				}
				
				if((tmpArray[5].equalsIgnoreCase("A") && tmpArray[6].equalsIgnoreCase("A")) || (tmpArray[5].equalsIgnoreCase("T") && tmpArray[6].equalsIgnoreCase("T")) ||  (tmpArray[5].equalsIgnoreCase("A") && tmpArray[6].equalsIgnoreCase("T")) || (tmpArray[5].equalsIgnoreCase("T") && tmpArray[6].equalsIgnoreCase("A"))){
					
				}
				else{
					if(!snpHomo){
						snpHeterLociC++;
					}
					else{
						snpHomoLociC++;
					}
				}
				
				CloseableIterator<SAMRecord> chrIt = inputSam.queryOverlapping(chr, snpPosition, snpPosition);
				record: while (chrIt.hasNext())
				{	
					SAMRecord samRecord = chrIt.next();
					int mapQual = samRecord.getMappingQuality();
					byte[] baseQual = samRecord.getBaseQualities();
					boolean unmapped = samRecord.getReadUnmappedFlag();
					if (unmapped || (mapQual < minMapQ))
					{
						continue record;
					}
					String seq = PicardUtils.getReadString(samRecord, true);
					try{
						
						String ref = PicardUtils.refStr(samRecord, true);
						boolean negStrand = samRecord.getReadNegativeStrandFlag();
						int alignmentS = samRecord.getAlignmentStart();
						//int readsStart = (negStrand) ? 0 - samRecord.getUnclippedEnd() : alignmentS;
						//int readsEnd = (negStrand) ? 0 - alignmentS : samRecord.getUnclippedEnd();
						int	onRefCoord = (negStrand) ? samRecord.getUnclippedEnd() : alignmentS; 
						int i = Math.abs(snpPosition - onRefCoord);
						if (seq.length() != ref.length())
						{
							System.err.println("SeqLen(" + seq.length() + ") != RefLen(" + ref.length() + ")");
							System.err.println(seq + "\n" + ref);
						}
						int seqLen = Math.min(seq.length(), ref.length());
						byte baseQS = (negStrand) ? baseQual[seqLen-1-i] : baseQual[i];
						if( baseQS < minBaseQual )
						{
							continue record;
						}
						if (ref.charAt(i) == '-')
						{
							break record;
						}
						
						totalNum++;
						if(PicardUtils.isCytosine(i, seq, true)){
							
							if(PicardUtils.isCytosine(i, seq, false)){
								totalNumC++;
							}
							else{
								totalNumT++;
							}
							
						}
						
					}
					catch (Exception e)
					{
						System.err.println("-----------------------------------------");
						System.err.println("Couldn't handle seq #");
						System.err.println(seq);
						e.printStackTrace(System.err);
						System.err.println("-----------------------------------------");
//						chrIt.close();
//						System.exit(1);
					}
					
				}
				chrIt.close();
				
				//push(@out,"$chr\t$start\t$end\t$name\t0\t+\n");
				if(totalNum >= minReadCov){
					outWriter.println(line);
					outBedWriter.println(chr + "\t" + start + "\t" + end + "\t" + name + "\t0\t+");
					if(snpHomo){
						snpHomoLociFilter++;
						
					}
					else{
						snpHeterLociFilter++;
						if((tmpArray[5].equalsIgnoreCase("C") && tmpArray[6].equalsIgnoreCase("T")) || (tmpArray[5].equalsIgnoreCase("T") && tmpArray[6].equalsIgnoreCase("C")) || (tmpArray[5].equalsIgnoreCase("A") && tmpArray[6].equalsIgnoreCase("G")) || (tmpArray[5].equalsIgnoreCase("G") && tmpArray[6].equalsIgnoreCase("A"))){
							ctSnpNumFilter++;
							outCTsnpWriter.println(line);
							outCTsnpBedWriter.println(chr + "\t" + start + "\t" + end + "\t" + name + "\t0\t+");
						}
						else{
							nonCtSnpNumFilter++;
							outNonCTsnpWriter.println(line);
							outNonCTsnpBedWriter.println(chr + "\t" + start + "\t" + end + "\t" + name + "\t0\t+");
						}
					}
					
					if((totalNumC + totalNumT) >= minCTReadCov){
						outCytosineWriter.println(line);
						outCytosineBedWriter.println(chr + "\t" + start + "\t" + end + "\t" + name + "\t0\t+");
						if((tmpArray[5].equalsIgnoreCase("A") && tmpArray[6].equalsIgnoreCase("A")) || (tmpArray[5].equalsIgnoreCase("T") && tmpArray[6].equalsIgnoreCase("T")) ||  (tmpArray[5].equalsIgnoreCase("A") && tmpArray[6].equalsIgnoreCase("T")) || (tmpArray[5].equalsIgnoreCase("T") && tmpArray[6].equalsIgnoreCase("A"))){
							
						}
						else{
							if(!snpHomo){
								snpHeterLociCFilter++;
							}
							else{
								snpHomoLociCFilter++;
							}
						}
					}
				}
				
				
					
			}
			inputSam.close();
			
			outWriter.close();
			outCytosineWriter.close();
			outBedWriter.close();
			outCytosineBedWriter.close();
			outCTsnpWriter.close();
			outCTsnpBedWriter.close();
			outNonCTsnpWriter.close();
			outNonCTsnpBedWriter.close();

			System.out.printf("heterozygous loci in 1M SNP array is: %d\n",snpHeterLoci);
			System.out.printf("homozygous loci in 1M SNP array is: %d\n",snpHomoLoci);
			System.out.printf("heterozygous CT SNP loci in 1M SNP array is: %d\n",ctSnpNum);
			System.out.printf("heterozygous non-CT SNP loci in 1M SNP array is: %d\n",nonCtSnpNum);
			
			
			System.out.printf("heterozygous loci in 1M SNP array after filter is: %d\n",snpHeterLociFilter);
			System.out.printf("homozygous loci in 1M SNP array after filter is: %d\n",snpHomoLociFilter);
			System.out.printf("heterozygous CT SNP loci in 1M SNP array after filter is: %d\n",ctSnpNumFilter);
			System.out.printf("heterozygous non-CT SNP loci in 1M SNP array after filter is: %d\n",nonCtSnpNumFilter);
			
			System.out.printf("heterozygous cytosine loci in 1M SNP array is: %d\n",snpHeterLociC);
			System.out.printf("homozygous cytosine loci in 1M SNP array is: %d\n",snpHomoLociC);

			System.out.printf("heterozygous cytosine loci in 1M SNP array after filter is: %d\n",snpHeterLociCFilter);
			System.out.printf("homozygous cytosine loci in 1M SNP array after filter is: %d\n",snpHomoLociCFilter);
			
			
			
		}

}
